"""
Origin_Image_Correction
包含两个函数对图像进行校正
"""

import cv2
import numpy as np
from math import sqrt


def line_points_function(point1, point2):
    """
    该函数的作用：计算两点之间的直线上所有点的位置坐标
    :param point1: 一点
    :param point2: 一点
    :return: 两点之间直线上所有点的坐标
    """
    k = (point1[1] - point2[1]) / (point1[0] - point2[0])
    b = point1[1] - k * point1[0]
    if point1[0] < point2[0]:
        x_start = point1[0]
        x_end = point2[0]
    else:
        x_start = point2[0]
        x_end = point1[0]
    line_range = [[x, int(k * x + b)] for x in range(x_start, x_end + 1)]
    return line_range


def line_exist(img, contours_points):
    """
    该函数的作用：根据两点之间直线的灰度值判断两点之间是否存在直线
    :param img: 原始图象
    :param contours_points: 初步寻找出来的表格框的四个角点
    :return: yes_flag : 返回合格
    """
    img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    x_lu = contours_points[0][0]
    y_lu = contours_points[0][1]

    x_ru = contours_points[1][0]
    y_ru = contours_points[1][1]

    x_ld = contours_points[2][0]
    y_ld = contours_points[2][1]

    x_rd = contours_points[3][0]
    y_rd = contours_points[3][1]

    points_test = [[x_lu, y_lu], [x_ru, y_ru], [x_ld, y_ld], [x_rd, y_rd]]
    line_flag = []
    for index in range(3):
        if index == 0:
            line_range = line_points_function(points_test[index], points_test[index + 1])
            sum_points = 0
            for i in range(len(line_range)):
                # print(img_gray[line_range[i][0]][line_range[i][1]])
                sum_points = img_gray[line_range[i][0]][line_range[i][1]] + sum_points
            line_flag.append(sum_points / len(line_range))
            line_range = line_points_function(points_test[index], points_test[index + 2])

            sum_points = 0
            for i in range(len(line_range)):
                # print(img_gray[line_range[i][0]][line_range[i][1]])
                sum_points = img_gray[line_range[i][0]][line_range[i][1]] + sum_points
            line_flag.append(sum_points / len(line_range))
        if index == 1:
            line_range = line_points_function(points_test[index], points_test[index + 2])
            sum_points = 0
            for i in range(len(line_range)):
                # print(img_gray[line_range[i][0]][line_range[i][1]])
                sum_points = img_gray[line_range[i][0]][line_range[i][1]] + sum_points
            line_flag.append(sum_points / len(line_range))
        if index == 2:
            line_range = line_points_function(points_test[index], points_test[index + 1])
            sum_points = 0
            for i in range(len(line_range)):
                # print(img_gray[line_range[i][0]][line_range[i][1]])
                sum_points = img_gray[line_range[i][0]][line_range[i][1]] + sum_points
            line_flag.append(sum_points / len(line_range))

    line_flag = np.array(line_flag)  # line_flag[左上与右上， 左上与左下， 右上与右下， 左下与右下]
    yes_flag = np.where(line_flag < 200)  # 将两点之间直线灰度值均值小于200即认为两点之间存在直线
    yes_flag = list(yes_flag)
    # 四个点的合格与否的标志值
    flag_0 = 0
    flag_1 = 0
    flag_2 = 0
    flag_3 = 0
    for flag in yes_flag[0]:
        if flag == 0:
            flag_0 = flag_0 + 1
            flag_1 = flag_1 + 1
        if flag == 1:
            flag_0 = flag_0 + 1
            flag_2 = flag_2 + 1
        if flag == 2:
            flag_1 = flag_1 + 1
            flag_3 = flag_3 + 1
        if flag == 3:
            flag_2 = flag_2 + 1
            flag_3 = flag_3 + 1
    # contours_point对应点的是否合格的标志值
    flag_list = [flag_0, flag_1, flag_2, flag_3]

    return flag_list


def find_jiaodian(xy_points, img):
    """
    该函数的作用：将图象角点检测得到的角点进行筛选得出初选的表格框的四个角点
    :param xy_points: 图象角点检测得到的角点
    :param img: 待识别图象
    :return: 返回表格框四个角点的坐标值
    """
    # 筛选轮廓四个角点
    # 其中x : row   y :col
    row_lenth = img.shape[0]
    col_lenth = img.shape[1]
    center_x = int(row_lenth / 2)
    center_y = int(col_lenth / 2)

    l_ru = 200
    l_lu = 200
    l_rd = 200
    l_ld = 200

    x_min = min(xy_points[0])
    x_max = max(xy_points[0])
    y_min = min(xy_points[1])
    y_max = max(xy_points[1])
    # print(x_min, x_max, y_min, y_max)

    x_left_up = center_x
    y_left_up = center_y
    x_right_up = center_x
    y_right_up = center_y
    x_left_d = center_x
    y_left_d = center_y
    x_right_down = center_x
    y_right_down = center_y

    for i in range(len(xy_points[0])):
        # print(contours[0][i][0][0],contours[0][i][0][1])
        x = xy_points[0][i] - x_min
        y = xy_points[1][i] - y_min
        l1 = sqrt(x ** 2 + y ** 2)
        if l1 < l_ru:
            #             x_ru_m = x
            #             y_ru_m = y
            l_ru = l1
            x_left_up = xy_points[0][i]
            y_left_up = xy_points[1][i]

        x = x_max - xy_points[0][i]
        y = xy_points[1][i] - y_min
        l2 = sqrt(x ** 2 + y ** 2)
        if l2 < l_lu:
            l_lu = l2
            x_left_d = xy_points[0][i]
            y_left_d = xy_points[1][i]

        x = xy_points[0][i] - x_min
        y = y_max - xy_points[1][i]
        l3 = sqrt(x ** 2 + y ** 2)
        if l3 < l_ld:
            l_ld = l3
            x_right_up = xy_points[0][i]
            y_right_up = xy_points[1][i]

        x = x_max - xy_points[0][i]
        y = y_max - xy_points[1][i]
        l4 = sqrt(x ** 2 + y ** 2)
        if l4 < l_rd:
            l_rd = l4
            x_right_down = xy_points[0][i]
            y_right_down = xy_points[1][i]
            # print(x_right_d)

    return [[x_left_up, y_left_up], [x_right_up, y_right_up], [x_left_d, y_left_d], [x_right_down, y_right_down]]


def contour_points_origin(img):
    img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    img_gray = cv2.dilate(img_gray, None)  # 膨胀：将前景物体变大，理解成将图像断开裂缝变小（在图片上画上黑色印记，印记越来越小）
    # 焦点检测必须是float32
    img_gray = np.float32(img_gray)
    # 角点检测
    dst = cv2.cornerHarris(img_gray, 2, 3, 0.04)  # 用来进行角点检测，数据类型为 float32 的输入图像。
    # dst = cv2.dilate(dst, (7, 7))
    # 角点检测可视化检验
    #     img_jiaodian = img
    #     img_jiaodian[dst>0.001*dst.max()]=[0,0,255]
    #     cv2.imwrite("C:\\Users\\lg\\Desktop\\form_ocr_excel\\hand write\\numbers_jiaodian.jpg", img_jiaodian)
    # np.where 返回一个元组，其中两个列表，分别为横纵坐标
    xy_points = np.where(dst > 0.0001 * dst.max())  # np.where(condition, x, y) 满足条件(condition)，输出x，不满足输出y。
    # print(dst > 0.001 * dst.max())
    # print(xy_points)
    # 筛选轮廓四个角点
    # 其中x : row   y :col
    l_ru = 50

    l_lu = 200

    l_rd = 200

    l_ld = 200

    x_min = min(xy_points[0])
    x_max = max(xy_points[0])
    y_min = min(xy_points[1])
    y_max = max(xy_points[1])
    # print(x_min, x_max, y_min, y_max)

    for i in range(len(xy_points[0])):
        # print(contours[0][i][0][0],contours[0][i][0][1])
        x = xy_points[0][i] - x_min
        y = xy_points[1][i] - y_min
        l1 = sqrt(x ** 2 + y ** 2)
        if (l1 < l_ru):
            #             x_ru_m = x
            #             y_ru_m = y
            l_ru = l1
            x_left_up = xy_points[0][i]
            y_left_up = xy_points[1][i]

        x = x_max - xy_points[0][i]
        y = xy_points[1][i] - y_min
        l2 = sqrt(x ** 2 + y ** 2)
        if (l2 < l_lu):
            l_lu = l2
            x_left_d = xy_points[0][i]
            y_left_d = xy_points[1][i]

        x = xy_points[0][i] - x_min
        y = y_max - xy_points[1][i]
        l3 = sqrt(x ** 2 + y ** 2)
        if (l3 < l_ld):
            l_ld = l3
            x_right_up = xy_points[0][i]
            y_right_up = xy_points[1][i]

        x = x_max - xy_points[0][i]
        y = y_max - xy_points[1][i]
        l4 = sqrt(x ** 2 + y ** 2)
        if (l4 < l_rd):
            l_rd = l4
            x_right_down = xy_points[0][i]
            y_right_down = xy_points[1][i]
            # print(x_right_d)

    return ([[x_left_up, y_left_up], [x_right_up, y_right_up], [x_left_d, y_left_d], [x_right_down, y_right_down]])


def contour_points_cornor(img):
    size = img.shape
    # print(size)
    y = size[0]  # 图片垂直尺寸高y
    x = size[1]  # 图片水平尺寸宽x

    # 一、初始化
    image_gauss = cv2.GaussianBlur(img, (3, 3), 0)  # Gauss滤波
    img_gray = cv2.cvtColor(image_gauss, cv2.COLOR_BGR2GRAY)  # 转化为灰度图
    ret, binary = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)  # 图像二值化

    # 二、轮廓提取
    contours, heriachy = cv2.findContours(~binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    area_max = 0
    for i, contour in enumerate(contours):  # 寻找最大轮廓
        #cv2.drawContours(img, contours, i, (0, 255, 0), 2)
        area = cv2.contourArea(contour)  # 用来计算轮廓的面积
        if area > area_max:
            area_max = area
            index = i
    #cv2.drawContours(img, contours, index, (0, 0, 255), 2)
    # print(index)
    # print (contours[index])

    #三、角点提取
    approxCurve = cv2.approxPolyDP(contours[index], 4, True)  # 用多边形包围轮廓，可以得到严格的矩形，有助于找到角点
    #print(approxCurve)

    points = approxCurve.flatten()  # 将矩阵变为一维列表
    # print(points)
    # print(int((len(points) + 1) / 2))
    d0, d1, d2, d3 = [], [], [], []
    for i in range(int((len(points) + 1) / 2)):
        d0.append((points[2 * i]) ** 2 + (points[2 * i + 1]) ** 2)
        d1.append((points[2 * i] - x) ** 2 + (points[2 * i + 1]) ** 2)
        d2.append((points[2 * i]) ** 2 + (points[2 * i + 1] - y) ** 2)
        d3.append((points[2 * i] - x) ** 2 + (points[2 * i + 1] - y) ** 2)

    p0 = (points[2 * d0.index(min(d0))], points[2 * d0.index(min(d0)) + 1])
    p1 = (points[2 * d1.index(min(d1))], points[2 * d1.index(min(d1)) + 1])
    p2 = (points[2 * d2.index(min(d2))], points[2 * d2.index(min(d2)) + 1])
    p3 = (points[2 * d3.index(min(d3))], points[2 * d3.index(min(d3)) + 1])
    #print(p0, p1, p2, p3)

    '''
    # 角点可视化
    cv2.circle(img, p0, 2, (0, 0, 255), 2)
    cv2.circle(img, p1, 2, (0, 0, 255), 2)
    cv2.circle(img, p2, 2, (0, 0, 255), 2)
    cv2.circle(img, p3, 2, (0, 0, 255), 2)
    # 将过大图片缩小两倍
    # img_origin = cv2.resize(img, (int(x / 2), int(y / 2)))
    cv2.imshow("binary2", img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    '''

    return ([[p0[1], p0[0]], [p1[1], p1[0]], [p2[1], p2[0]], [p3[1], p3[0]]])

def contour_points(img):
    """
    该函数的作用：寻找待检测图象最终的表格框的四个角点
    :param img: 待检测图象
    :return: 最终检测得到的图象中表格区域的四个角点坐标值
    """

    img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    img_gray = cv2.dilate(img_gray, None)
    # 焦点检测必须是float32
    img_gray = np.float32(img_gray)
    # 角点检测
    dst = cv2.cornerHarris(img_gray, 2, 3, 0.04)
    # dst = cv2.dilate(dst, (7, 7))
    # 角点检测可视化检验
    #     img_jiaodian = img
    #     img_jiaodian[dst>0.001*dst.max()]=[0,0,255]
    #     cv2.imwrite("H:\\20201111_dqform\\202115excel_roi\\numbers_jiaodian.jpg", img_jiaodian)
    # np.where 返回一个元组，其中两个列表，分别为横纵坐标
    xy_points = np.where(dst > 0.001 * dst.max())
    xy_points = list(xy_points)
    # 筛选角点，去掉一些靠近图像中心，离边框较远的点
    row_lenth = img.shape[0]
    col_lenth = img.shape[1]
    center_x = int(row_lenth / 2)
    center_y = int(col_lenth / 2)
    l_x = []
    l_y = []
    for i in range(len(xy_points[0])):
        row_x = xy_points[0][i]
        col_y = xy_points[1][i]
        if abs(row_x - center_x) > 0.8 * center_x or abs(col_y - center_y) > 0.8 * center_y:
            l_x.append(xy_points[0][i])
            l_y.append(xy_points[1][i])

    xy_points = [l_x, l_y]
    while 1:

        contours_points = find_jiaodian(xy_points, img)
        flag_list = line_exist(img, contours_points)
        unqualified_points_index = []  # contours_points中不合格角点的索引
        for index_flag in range(len(flag_list)):
            if flag_list[index_flag] == 0:
                unqualified_points_index.append(index_flag)
        if len(unqualified_points_index) == 0:  # 如果不合格角点索引为空即所用角点均合格，因此退出循环！
            break
        # 将不合格的角点从角点列表中剔除，进入下一轮循环
        for index_points in unqualified_points_index:
            for index_point_find in range(len(xy_points[0])):
                if xy_points[1][index_point_find] == contours_points[index_points][1] and xy_points[0][
                    index_point_find] == contours_points[index_points][0]:
                    xy_points[0].pop(index_point_find)
                    xy_points[1].pop(index_point_find)
                    break

    return contours_points


# 校正图像
# 利用均值来对图像进行校正
def img_crroect(jiaozheng_img_path, contours_points, img):
    """
    该函数的作用：利用四个角点均值之间的关系，通过透射变换对图像进行矫正
    :param jiaozheng_img_path:
    :param contours_points:
    :param img:
    :return:
    """
    # pts1 = np.float32( [[x_right_up, y_right_up],[x_left_d, y_left_d],[x_right_d, y_right_d]])

    # pts2 = np.float32([[x_left_up, y_right_up],[x_left_d, y_left_up],[x_left_d, y_right_up]])
    rows, cols, ch = img.shape
    x1 = int((contours_points[0][0] + contours_points[1][0]) / 2)
    y1 = int((contours_points[1][1] + contours_points[3][1]) / 2)

    x2 = int((contours_points[2][0] + contours_points[3][0]) / 2)
    y2 = int((contours_points[0][1] + contours_points[2][1]) / 2)

    #     pts1 = np.float32([contours_points[1], contours_points[2],contours_points[3]])
    #     pts2 = np.float32([ [x1, y1], [x2, y2], [x2, y1]])

    # 这里注意一下：透射变换的坐标要求是w,h格式，而contour_points里面点的坐标格式是h, w
    pts1 = np.float32([[contours_points[0][1], contours_points[0][0]], [contours_points[1][1], contours_points[1][0]],
                       [contours_points[2][1], contours_points[2][0]], [contours_points[3][1], contours_points[3][0]]])

    # 利用框线的四个角点均值对框线进行校正
    # 左上角点
    w1 = int((contours_points[0][1] + contours_points[2][1]) / 2)
    h1 = int((contours_points[0][0] + contours_points[1][0]) / 2)

    # 右上角点
    w2 = int((contours_points[1][1] + contours_points[3][1]) / 2)
    h2 = h1

    # 左下角点
    w3 = w1
    h3 = int((contours_points[2][0] + contours_points[3][0]) / 2)

    # 右下角点
    w4 = w2
    h4 = h3

    pts2 = np.float32([[contours_points[0][1], contours_points[0][0]], [contours_points[1][1], contours_points[0][1]],
                       [contours_points[0][1], contours_points[2][0]], [contours_points[1][1], contours_points[2][0]]])

    # 利用四个角点的均值作为透射变换求解变换矩阵的值
    pts2 = np.float32([[w1, h1], [w2, h2],
                       [w3, h3], [w4, h4]])

    # 获得放射矩阵
    #     M = cv2.getAffineTransform(pts2, pts1)
    #     # 利用仿射矩阵对图像进行旋转
    #     dst = cv2.warpAffine(img, M, (cols, rows))
    M = cv2.getPerspectiveTransform(pts1, pts2)
    dst = cv2.warpPerspective(img, M, (cols, rows))
    # 对旋转之后的图像进行简单的剪裁，除去一些无谓的计算
    # dst = dst[h1 - 5: h3 + 5, w1 - 5: w2 + 5, :]
    dst = dst[h1 : h3 , w1 : w2 , :]



    # plt.imshow(dst)
    if (cv2.imwrite(jiaozheng_img_path, dst)):
        print("图像校正成功，并成功保存！")
    else:
        print("图像校正失败！！！")
    return dst
